import contextlib
from functools import partial

import torch

@contextlib.contextmanager
def capture_modules_output(modules: dict[str, torch.nn.Module]):
    modules_output = {}
    modules_handle = set()

    def hook(name, module, input, output):
        modules_output[name] = output

    for name, module in modules.items():
        handle = module.register_forward_hook(partial(hook, name))
        modules_handle.add(handle)

    try:
        yield modules_output
    finally:
        for handle in modules_handle:
            handle.remove()
